import asyncio
from src.module.Model import Model
from langchain.prompts import ChatPromptTemplate

from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from langchain.globals import set_verbose, set_debug

# 开启 verbose 模式, verbose 只会输出重要的事件
set_verbose(True)

# 开启 debug 模式，debug 会输出所有的事件
# set_debug(True)

class OutputParser(Model):
    @staticmethod
    async def async_handler(chain):
        async for chunk in chain.astream({"question": [HumanMessage(content='hello world!')]}):
            print(chunk)

    @staticmethod
    async def json_output(chain):
        async for chunk in chain.astream({"question": [HumanMessage(content='today is a good day!')]}):
            print(chunk)

    def start(self):


        # 输出为 str 格式
        # template = ChatPromptTemplate.from_messages([
        #     ('system', '请将下面的内容翻译为中文'),
        #     ('placeholder', '{question}')
        # ])
        # parser = StrOutputParser()
        # 链式调用，第一步、第二步、第三步。每一步必须要有 invoke 可执行函数
        # chain = template | self.model | parser
        # 同步输出
        # print(chain.invoke({"question": [HumanMessage(content='hello world!')]}))
        # stream 输出
        # full_chunk = ''
        # for chunk in chain.stream({"question": [HumanMessage(content='hello world!')]}):
        #     full_chunk += chunk
        #     print(chunk, end='|', flush=True)
        # 协程的异步输出
        # asyncio.run(self.async_handler(chain))

        # 输出为 json 格式
        # 需要注意的是，一定要在问题中规定回答的方式是json，否则， chunk 将不会有任何输出
        prompt = ChatPromptTemplate.from_messages([
            ('system', '请将下面的内容翻译为中文, 以 JSON 的方式返回给我'),
            ('placeholder', '{question}')
        ])
        parser = JsonOutputParser()
        chain = prompt | self.model | parser
        asyncio.run(self.json_output(chain))

